package com.jd.finance.util;

import com.alibaba.excel.EasyExcel;
import com.alibaba.excel.ExcelWriter;
import com.alibaba.excel.write.metadata.WriteSheet;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.net.URLEncoder;
import java.util.Collection;
import java.util.Collections;

/**
 * Excel Util
 */
public class EasyExcelUtil {

    /**
     * Export Data to Excel.
     *
     * @param <T>
     * @param fileName  Excel file name.
     * @param sheetName Excel sheet name.
     * @param cls       data class
     * @param data      page data.
     * @throws IOException IOException
     */
    public static <T> void export(String fileName, String sheetName, Class<T> cls, Collection<?> data)
            throws IOException {
        ServletRequestAttributes servletContainer = (ServletRequestAttributes) RequestContextHolder
                .getRequestAttributes();
        HttpServletResponse response = servletContainer.getResponse();
        response.setContentType("application/vnd.openxmlformats-officedocument.spreadsheetml.sheet");
        response.setCharacterEncoding("utf-8");
        String fname = URLEncoder.encode(fileName, "UTF-8");
        response.setHeader("Content-disposition", "attachment;filename=" + fname + ".xlsx");
        EasyExcel.write(response.getOutputStream(), cls).sheet(sheetName).doWrite(Collections.singletonList(data));
    }

    /**
     * Export Data to Excel and which includes two sheets.
     *
     * @param <T>         the first sheet Class.
     * @param <K>         the second sheet Class.
     * @param fileName    Excel file name.
     * @param sheet1Name  the first sheet name.
     * @param sheet2Name  the second sheet name.
     * @param sheet1Class the first sheet Class.
     * @param sheet2Class the second sheet Class.
     * @param sheet1Datas the first sheet data.
     * @param sheet2Datas the second sheet data.
     * @throws IOException Exception
     */
    public static <T, K> void export2Sheet(String fileName, String sheet1Name, String sheet2Name, Class<T> sheet1Class,
                                           Class<K> sheet2Class, Collection<?> sheet1Datas, Collection<?> sheet2Datas) throws IOException {
        ServletRequestAttributes servletContainer = (ServletRequestAttributes) RequestContextHolder
                .getRequestAttributes();
        HttpServletResponse response = servletContainer.getResponse();
        response.setContentType("application/vnd.openxmlformats-officedocument.spreadsheetml.sheet");
        response.setCharacterEncoding("utf-8");
        String fname = URLEncoder.encode(fileName, "UTF-8");
        response.setHeader("Content-disposition", "attachment;filename=" + fname + ".xlsx");
        ExcelWriter excelWriter = EasyExcel.write(response.getOutputStream()).build();
        WriteSheet firstSheet = EasyExcel.writerSheet(0, sheet1Name).head(sheet1Class).build();
        excelWriter.write(Collections.singletonList(sheet1Datas), firstSheet);
        WriteSheet secondSheet = EasyExcel.writerSheet(1, sheet2Name).head(sheet2Class).build();
        excelWriter.write(Collections.singletonList(sheet2Datas), secondSheet);
        excelWriter.finish();
    }

}
